"""
MCP Server 实现
实现工具、资源和提示的核心功能
"""
import asyncio
import json
import os
from datetime import datetime
from typing import Any, Dict, List, Optional
from pathlib import Path

from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import aiofiles

from mcp_models import (
    Tool, Resource, Prompt, ToolCall, ToolResult, 
    ResourceRequest, ResourceResponse, PromptRequest, PromptResponse,
    MCPCapabilities, ToolParameter, ResourceType
)


class MCPServer:
    """MCP服务器核心实现"""
    
    def __init__(self):
        self.app = FastAPI(title="MCP Server", version="1.0.0")
        self.setup_middleware()
        self.setup_routes()
        
        # 存储注册的工具、资源和提示
        self.tools: Dict[str, Tool] = {}
        self.resources: Dict[str, Resource] = {}
        self.prompts: Dict[str, Prompt] = {}
        
        # 连接的客户端
        self.connected_clients: List[WebSocket] = []
        
        # 初始化默认能力
        self._register_default_capabilities()
    
    def setup_middleware(self):
        """设置中间件"""
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
    
    def setup_routes(self):
        """设置路由"""
        
        # 静态文件服务
        self.app.mount("/static", StaticFiles(directory="static"), name="static")
        
        @self.app.get("/")
        async def root():
            return FileResponse("static/index.html")
        
        @self.app.get("/api/status")
        async def status():
            return {"message": "MCP Server is running", "version": "1.0.0"}
        
        @self.app.get("/capabilities", response_model=MCPCapabilities)
        async def get_capabilities():
            """获取服务器能力"""
            return MCPCapabilities(
                tools=list(self.tools.values()),
                resources=list(self.resources.values()),
                prompts=list(self.prompts.values()),
                server_info={
                    "name": "Demo MCP Server",
                    "version": "1.0.0",
                    "description": "演示MCP协议的三个核心基元"
                }
            )
        
        @self.app.post("/tools/call", response_model=ToolResult)
        async def call_tool(tool_call: ToolCall):
            """调用工具"""
            return await self.execute_tool(tool_call)
        
        @self.app.post("/resources/get", response_model=ResourceResponse)
        async def get_resource(resource_request: ResourceRequest):
            """获取资源"""
            return await self.fetch_resource(resource_request)
        
        @self.app.post("/prompts/render", response_model=PromptResponse)
        async def render_prompt(prompt_request: PromptRequest):
            """渲染提示"""
            return await self.render_prompt_template(prompt_request)
        
        @self.app.websocket("/ws")
        async def websocket_endpoint(websocket: WebSocket):
            """WebSocket连接处理"""
            await self.handle_websocket(websocket)
    
    def _register_default_capabilities(self):
        """注册默认的工具、资源和提示"""
        
        # 注册工具 - 模型控制
        self.register_tool(Tool(
            name="send_notification",
            description="发送系统通知",
            parameters=[
                ToolParameter(name="message", type="string", description="通知内容", required=True),
                ToolParameter(name="priority", type="string", description="优先级", required=False, default="normal")
            ]
        ))
        
        self.register_tool(Tool(
            name="create_file",
            description="创建文件",
            parameters=[
                ToolParameter(name="filename", type="string", description="文件名", required=True),
                ToolParameter(name="content", type="string", description="文件内容", required=True)
            ]
        ))
        
        self.register_tool(Tool(
            name="execute_command",
            description="执行系统命令",
            parameters=[
                ToolParameter(name="command", type="string", description="要执行的命令", required=True),
                ToolParameter(name="timeout", type="integer", description="超时时间(秒)", required=False, default=30)
            ]
        ))
        
        # 注册资源 - 应用控制
        self.register_resource(Resource(
            uri="file://logs/{date}.log",
            name="system_logs",
            description="系统日志文件",
            type=ResourceType.TEXT,
            template_params={"date": "YYYY-MM-DD格式的日期"}
        ))
        
        self.register_resource(Resource(
            uri="api://weather/{city}",
            name="weather_info",
            description="天气信息",
            type=ResourceType.JSON,
            template_params={"city": "城市名称"}
        ))
        
        # 注册提示 - 用户控制
        self.register_prompt(Prompt(
            name="code_review",
            description="代码审查模板",
            template="请对以下代码进行审查：\n\n```{language}\n{code}\n```\n\n审查重点：{focus_areas}\n\n请提供详细的审查意见和改进建议。",
            parameters=[
                ToolParameter(name="code", type="string", description="要审查的代码", required=True),
                ToolParameter(name="language", type="string", description="编程语言", required=False, default="python"),
                ToolParameter(name="focus_areas", type="string", description="审查重点", required=False, default="安全性、性能、可维护性")
            ]
        ))
        
        self.register_prompt(Prompt(
            name="documentation",
            description="文档生成模板",
            template="为以下{type}生成详细文档：\n\n{content}\n\n要求：\n1. 包含功能描述\n2. 参数说明\n3. 使用示例\n4. 注意事项",
            parameters=[
                ToolParameter(name="content", type="string", description="要生成文档的内容", required=True),
                ToolParameter(name="type", type="string", description="内容类型", required=False, default="函数")
            ]
        ))
    
    def register_tool(self, tool: Tool):
        """注册工具"""
        self.tools[tool.name] = tool
    
    def register_resource(self, resource: Resource):
        """注册资源"""
        self.resources[resource.name] = resource
    
    def register_prompt(self, prompt: Prompt):
        """注册提示"""
        self.prompts[prompt.name] = prompt
    
    async def execute_tool(self, tool_call: ToolCall) -> ToolResult:
        """执行工具调用"""
        if tool_call.tool_name not in self.tools:
            return ToolResult(
                request_id=tool_call.request_id,
                success=False,
                error=f"Tool '{tool_call.tool_name}' not found"
            )
        
        try:
            if tool_call.tool_name == "send_notification":
                result = await self._send_notification(tool_call.parameters)
            elif tool_call.tool_name == "create_file":
                result = await self._create_file(tool_call.parameters)
            elif tool_call.tool_name == "execute_command":
                result = await self._execute_command(tool_call.parameters)
            else:
                result = {"message": f"Tool '{tool_call.tool_name}' executed successfully", "parameters": tool_call.parameters}
            
            return ToolResult(
                request_id=tool_call.request_id,
                success=True,
                result=result
            )
        except Exception as e:
            return ToolResult(
                request_id=tool_call.request_id,
                success=False,
                error=str(e)
            )
    
    async def fetch_resource(self, resource_request: ResourceRequest) -> ResourceResponse:
        """获取资源"""
        # 简化的资源获取逻辑
        uri = resource_request.uri
        
        if uri.startswith("file://"):
            return await self._fetch_file_resource(uri, resource_request.parameters)
        elif uri.startswith("api://weather/"):
            return await self._fetch_weather_resource(uri, resource_request.parameters)
        else:
            raise HTTPException(status_code=404, detail=f"Resource not found: {uri}")
    
    async def render_prompt_template(self, prompt_request: PromptRequest) -> PromptResponse:
        """渲染提示模板"""
        if prompt_request.name not in self.prompts:
            raise HTTPException(status_code=404, detail=f"Prompt '{prompt_request.name}' not found")
        
        prompt = self.prompts[prompt_request.name]
        
        try:
            rendered = prompt.template.format(**prompt_request.parameters)
            return PromptResponse(
                name=prompt_request.name,
                rendered_template=rendered,
                metadata={"rendered_at": datetime.now().isoformat()}
            )
        except KeyError as e:
            raise HTTPException(status_code=400, detail=f"Missing parameter: {e}")
    
    async def _send_notification(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """发送通知工具实现"""
        message = params.get("message", "")
        priority = params.get("priority", "normal")
        
        # 模拟发送通知
        notification = {
            "id": f"notif_{datetime.now().timestamp()}",
            "message": message,
            "priority": priority,
            "timestamp": datetime.now().isoformat(),
            "sent": True
        }
        
        # 通过WebSocket向所有连接的客户端发送通知
        if self.connected_clients:
            notification_json = json.dumps({"type": "notification", "data": notification})
            for client in self.connected_clients[:]:  # 创建副本以避免修改时迭代
                try:
                    await client.send_text(notification_json)
                except:
                    self.connected_clients.remove(client)
        
        return notification
    
    async def _create_file(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """创建文件工具实现"""
        filename = params.get("filename")
        content = params.get("content", "")
        
        if not filename:
            raise ValueError("filename is required")
        
        # 确保文件在安全目录内
        safe_path = Path("temp_files") / filename
        safe_path.parent.mkdir(parents=True, exist_ok=True)
        
        async with aiofiles.open(safe_path, 'w', encoding='utf-8') as f:
            await f.write(content)
        
        return {
            "filename": filename,
            "path": str(safe_path),
            "size": len(content.encode('utf-8')),
            "created_at": datetime.now().isoformat()
        }
    
    async def _execute_command(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """执行命令工具实现（安全性受限）"""
        command = params.get("command", "")
        timeout = params.get("timeout", 30)
        
        # 安全性：只允许特定的安全命令
        allowed_commands = ["echo", "date", "pwd", "ls", "dir"]
        cmd_parts = command.split()
        
        if not cmd_parts or cmd_parts[0] not in allowed_commands:
            raise ValueError(f"Command not allowed. Allowed commands: {allowed_commands}")
        
        try:
            process = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            
            stdout, stderr = await asyncio.wait_for(
                process.communicate(), 
                timeout=timeout
            )
            
            return {
                "command": command,
                "exit_code": process.returncode,
                "stdout": stdout.decode('utf-8', errors='ignore'),
                "stderr": stderr.decode('utf-8', errors='ignore'),
                "executed_at": datetime.now().isoformat()
            }
        except asyncio.TimeoutError:
            raise ValueError(f"Command timed out after {timeout} seconds")
    
    async def _fetch_file_resource(self, uri: str, params: Optional[Dict[str, Any]]) -> ResourceResponse:
        """获取文件资源"""
        # 解析文件路径
        file_path = uri.replace("file://", "")
        
        # 处理模板参数
        if params:
            for key, value in params.items():
                file_path = file_path.replace(f"{{{key}}}", str(value))
        
        # 安全性检查
        safe_path = Path("temp_files") / file_path
        
        if not safe_path.exists():
            # 创建示例日志文件
            safe_path.parent.mkdir(parents=True, exist_ok=True)
            sample_content = f"[{datetime.now().isoformat()}] INFO: Sample log entry\n[{datetime.now().isoformat()}] DEBUG: System initialized\n"
            async with aiofiles.open(safe_path, 'w') as f:
                await f.write(sample_content)
        
        async with aiofiles.open(safe_path, 'r', encoding='utf-8') as f:
            content = await f.read()
        
        return ResourceResponse(
            uri=uri,
            content=content,
            content_type="text/plain",
            metadata={"size": len(content), "accessed_at": datetime.now().isoformat()}
        )
    
    async def _fetch_weather_resource(self, uri: str, params: Optional[Dict[str, Any]]) -> ResourceResponse:
        """获取天气资源（模拟）"""
        city = uri.split("/")[-1]
        
        # 模拟天气数据
        weather_data = {
            "city": city,
            "temperature": 22,
            "humidity": 65,
            "weather": "晴朗",
            "wind_speed": 5.2,
            "timestamp": datetime.now().isoformat()
        }
        
        return ResourceResponse(
            uri=uri,
            content=json.dumps(weather_data, ensure_ascii=False, indent=2),
            content_type="application/json",
            metadata={"source": "simulated", "city": city}
        )
    
    async def handle_websocket(self, websocket: WebSocket):
        """处理WebSocket连接"""
        await websocket.accept()
        self.connected_clients.append(websocket)
        
        try:
            # 发送欢迎消息
            await websocket.send_text(json.dumps({
                "type": "welcome",
                "message": "Connected to MCP Server",
                "capabilities": {
                    "tools": len(self.tools),
                    "resources": len(self.resources),
                    "prompts": len(self.prompts)
                }
            }))
            
            while True:
                data = await websocket.receive_text()
                message = json.loads(data)
                
                # 处理客户端消息
                if message.get("type") == "ping":
                    await websocket.send_text(json.dumps({"type": "pong"}))
                
        except WebSocketDisconnect:
            pass
        finally:
            if websocket in self.connected_clients:
                self.connected_clients.remove(websocket)


# 创建服务器实例
mcp_server = MCPServer()
app = mcp_server.app

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000) 